from PySide6.QtWidgets import (
    QWidget, QFormLayout, QComboBox, QCheckBox, QLabel
)
from PySide6.QtCore import Signal
import os

from src.config import app_config
from src.ui.i18n import _


class GenerationOptionsWidget(QWidget):
    WHISPER_LANGUAGES = {
        "auto": "Auto Detect",
        "en": "English",
        "zh": "Chinese",
        "ja": "Japanese",
        "ko": "Korean",
        "fr": "French",
        "de": "German",
        "es": "Spanish",
    }

    def __init__(self, parent=None):
        super().__init__(parent)
        self._setup_ui()
        self._connect_signals()
        self.populate_models()
        self.retranslate_ui()

    def _setup_ui(self):
        layout = QFormLayout(self)
        layout.setContentsMargins(0, 0, 0, 0)

        self.model_combo = QComboBox()
        self.language_combo = QComboBox()
        for code, text_key in self.WHISPER_LANGUAGES.items():
            self.language_combo.addItem(text_key, code)

        self.translate_to_eng_checkbox = QCheckBox()

        self.whisper_model_label = QLabel()
        self.video_lang_label = QLabel()

        layout.addRow(self.whisper_model_label, self.model_combo)
        layout.addRow(self.video_lang_label, self.language_combo)
        layout.addRow(self.translate_to_eng_checkbox)

    def _connect_signals(self):
        self.model_combo.currentIndexChanged.connect(self._on_model_changed)

    def retranslate_ui(self):
        self.whisper_model_label.setText(_("Whisper Model:"))
        self.video_lang_label.setText(_("Video Language:"))
        self.translate_to_eng_checkbox.setText(_("Translate to English as well"))

        for i in range(self.language_combo.count()):
            code = self.language_combo.itemData(i)
            if code in self.WHISPER_LANGUAGES:
                self.language_combo.setItemText(i, _(self.WHISPER_LANGUAGES[code]))

    def populate_models(self):
        self.model_combo.blockSignals(True)
        self.model_combo.clear()

        # 1. Add local models
        models_dir = app_config.get('paths.models', 'models')
        if not os.path.isdir(models_dir):
            os.makedirs(models_dir)
        local_models = [f for f in os.listdir(models_dir) if f.endswith(".bin")]
        if local_models:
            for model in local_models:
                self.model_combo.addItem(model, userData={'type': 'local', 'name': model})

        # 2. Add online models from service providers
        providers = app_config.get("service_providers", [])
        online_models_added = False
        for provider in providers:
            if provider.get("enabled"):
                stt_service = provider.get("services", {}).get("stt", {})
                models = stt_service.get("models", [])
                if models:
                    if not online_models_added:
                        self.model_combo.addItem(_("--- Online Models ---"), userData=None)
                        online_models_added = True
                    for model in models:
                        display_name = f"{provider.get('provider_name')} / {model.get('display_name')}"
                        self.model_combo.addItem(display_name, userData={'type': 'online', 'name': model.get('model_name')})
        
        self._restore_last_used_model()
        self.model_combo.blockSignals(False)

    def _restore_last_used_model(self):
        last_used_model = app_config.get("generation.last_used_whisper_model")
        if last_used_model:
            last_used_model_name = last_used_model if isinstance(last_used_model, str) else last_used_model.get('name')
            if last_used_model_name:
                for i in range(self.model_combo.count()):
                    user_data = self.model_combo.itemData(i)
                    if user_data and user_data.get('name') == last_used_model_name:
                        self.model_combo.setCurrentIndex(i)
                        break

    def _on_model_changed(self, index):
        model_data = self.model_combo.itemData(index)
        if model_data:
            app_config.set("generation.last_used_whisper_model", model_data, emit_signal=False)
            app_config.save(emit_signal=False)

    def get_options(self):
        return {
            "whisper_model": self.model_combo.currentData(),
            "video_lang": self.language_combo.currentData(),
            "translate_to_english": self.translate_to_eng_checkbox.isChecked(),
        }